"""
《邢不行-2023新版|Python股票量化投资课程》
author: 邢不行
微信: xbx9585

根据选股数据，进行选股
"""
from Evaluate import *
from Filter import *
from Functions import *
import warnings
from Config import *
import datetime

warnings.filterwarnings('ignore')

pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
pd.set_option('display.max_rows', 5000)  # 最多显示数据的行数

print('策略名称:', strategy_name)
print('周期:', period_type)

# ===导入数据
# 从pickle文件中读取整理好的所有股票数据
df = pd.read_pickle(root_path + '/data/equity_curve/每周涨跌幅_%s_%s_%s.pkl' % (strategy_name, select_stock_num, period_type))

index_data = import_index_data(root_path + '/data/index_data/sh000300.csv')
df = pd.merge(left=df, right=index_data, on='交易日期', how='right', sort=True, indicator=True)
# 使用ffill方法填充缺失的值
df['股票数量'] = df['股票数量'].ffill()
df['买入股票代码'] = df['买入股票代码'].ffill()
df['买入股票名称'] = df['买入股票名称'].ffill()
df['选股下周期涨跌幅'] = df['选股下周期涨跌幅'].ffill()
df['选股下周期每天涨跌幅'] = df['选股下周期每天涨跌幅'].ffill()
df['资金曲线'] = df['资金曲线'].ffill()

df['股票数量'] = df['股票数量'].shift()
df['买入股票代码'] = df['买入股票代码'].shift()
df['买入股票名称'] = df['买入股票名称'].shift()
df['选股下周期涨跌幅'] = df['选股下周期涨跌幅'].shift()
df['选股下周期每天涨跌幅'] = df['选股下周期每天涨跌幅'].shift()
df['资金曲线'] = df['资金曲线'].shift()

df.dropna(subset=['指数涨跌幅'], inplace=True)
df.dropna(subset=['股票数量'], inplace=True)
df['stock_count'] = df.groupby('资金曲线').cumcount() + 1
df['stock_count'] = df['stock_count']-1

df['选股涨跌幅'] = df.apply(lambda row: row['选股下周期每天涨跌幅'][row['stock_count']] if isinstance(row['选股下周期每天涨跌幅'], list) and len(row['选股下周期每天涨跌幅']) > row['stock_count'] else None, axis=1)


df['资金曲线2'] = (df['选股涨跌幅'] + 1).cumprod()
df['开盘价']=df['资金曲线2']
df['最高价']=df['资金曲线2']
df['收盘价']=df['资金曲线2']
df['最低价']=df['资金曲线2']
df['涨跌幅']=df['选股涨跌幅']
df['持有股票代码']=df['买入股票代码']
df=df[['交易日期','指数涨跌幅','持有股票代码','涨跌幅','开盘价','最高价','收盘价','最低价']]
df.to_csv(root_path + '/data/equity_curve/a.csv',encoding='gbk')
df.to_pickle(root_path + '/data/equity_curve/每日资金曲线_%s_%s_%s.pkl' % (strategy_name, select_stock_num, period_type))
# 打印结果
print(df)